﻿using System;
using System.IO;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

using Z3 = Microsoft.Z3;

using System.Text.RegularExpressions;

namespace BMC
{
    class Program
    {
        static InputFile inputFile;
        static int searchBound;

        static void Main(string[] args)
        {
            if (args.Length != 2)
            {
                Console.WriteLine("Wrong number of arguments. Quitting.\n");
                return;
            }
            else if (System.IO.File.Exists(args[0]) == false)
            {
                Console.WriteLine("File does not exist. Quitting.\n");
                return;
            }
            else if (!int.TryParse(args[1], out searchBound))
            {
                Console.WriteLine("Search bound is not an integer. Quitting.\n");
                return;
            }
            List<String> _out = new List<string>();
            try
            {
                //Read input file
                String[] programStr = System.IO.File.ReadAllLines(args[0]);
                inputFile = new InputFile(programStr);

                
                //Redirect output
                FileStream ostrm;
                StreamWriter writer;
                TextWriter oldOut = Console.Out;
                try
                {
                    ostrm = new FileStream("./consoleOut.txt", FileMode.Create, FileAccess.Write);
                    writer = new StreamWriter(ostrm);
                }
                catch (Exception e)
                {
                    Console.WriteLine("Cannot open consoleOut.txt for writing. Quitting.");
                    Console.WriteLine(e.Message);
                    return;
                }
                Console.SetOut(writer);

                Console.WriteLine("Software Model Checking - Assignment 1, tasks 2 and 3");
                Console.WriteLine("Authors: Roman Kern, Alexander Prostejovsky\r\n");
                Console.WriteLine("*** Input program:");
                Console.WriteLine(String.Join("\r\n", inputFile.program)+"\r\n");

                Console.WriteLine("*** Parsed SMT2 program:");
                Smt2Program smt2program = new Smt2Program(inputFile,searchBound);
                _out = new List<String>(smt2program.smt2variables);
                _out.AddRange(smt2program.smt2inputs);
                _out.AddRange(smt2program.smt2initializations);
                _out.AddRange(smt2program.smt2statements);

                String prog = "(set-logic QF_BV)\n" + String.Join("\r\n", _out);
                
                Dictionary<string, string> cfg = new Dictionary<string, string>() { 
                { "MODEL", "true" },
                };
                using (Z3.Context ctx = new Z3.Context(cfg))
                {
                    // Use a solver for QF_BV
                    Z3.Solver s = ctx.MkSolver();

                    smt2program.smt2conditions.ForEach(
                        a =>
                            {
                                s.Push();
                                String prog2 = prog + "\r\n" + a;
                                Z3.BoolExpr assertions = ctx.ParseSMTLIB2String(prog2);
                                s.Assert(assertions);
                                Console.WriteLine("\r\n*** Checking\r\n" + a);
                                Z3.Status res = s.Check();
                                Console.WriteLine("Solver result: " + res);
                                if (res.Equals(Z3.Status.SATISFIABLE))
                                {
                                    Console.WriteLine("Counterexample trace:");
                                    Console.WriteLine(s.Model);
                                }
                                else
                                {
                                    Console.WriteLine("I'm sorry Dave, I'm afraid I can't do that.");
                                }
                                s.Pop();
                            }
                        );
                }

                Console.SetOut(oldOut);
                writer.Close();
                ostrm.Close();
            }
            catch (Exception e)
            {
                Console.WriteLine(e.Message);
                Console.WriteLine(e.StackTrace);
            }
        }

    }

    class Smt2Program
    {
        public Smt2Program(InputFile inputFile, int searchBound)
        {
            this.inputFile = inputFile;
            this.searchBound = searchBound;
            this.smt2variables = new List<String>();
            this.smt2inputs = new List<String>();
            this.smt2initializations = new List<String>();
            this.smt2statements = new List<String>();
            this.smt2conditions = new List<String>();
            this.arrayVariables = new List<ArrayVariable>();
            this.progSteps = 0;

            translateVariables();
            translateInputs();
            translateInitializations();
            translateStatement();
            translateConditions();
        }

        public InputFile inputFile
        {
            get;
            private set;
        }

        private int searchBound
        {
            get;
            set;
        }

        public List<String> smt2variables
        {
            get;
            private set;
        }

        public List<String> smt2inputs
        {
            get;
            private set;
        }

        public List<String> smt2initializations
        {
            get;
            private set;
        }

        public List<String> smt2statements
        {
            get;
            private set;
        }

        public List<String> smt2conditions
        {
            get;
            private set;
        }

        private List<ArrayVariable> arrayVariables
        {
            get;
            set;
        }

        private int progSteps
        {
            get;
            set;
        }

        private void translateVariables()
        {
            foreach (String varStr in inputFile.variables())
            {
                String[] splitStr = varStr.Split(' ');
                String act = splitStr[0];
                String next = splitStr[1];
                String type = splitStr[3];
                String bitwidth = splitStr[4].Substring(0, splitStr[4].Count() - 1);

                ArrayVariable var = new ArrayVariable(act, next, type, int.Parse(bitwidth), inputFile.statements().Count());
                arrayVariables.Add(var);
                //Console.WriteLine(var.ToString());
            }

            foreach(ArrayVariable var in arrayVariables)
            {
                smt2variables.Add("(declare-const " + var.act + " (Array Int (_ " + var.type + " " + var.bitwidth.ToString() + ")))");
            }

            foreach (String str in smt2variables)
            {
                Console.WriteLine(str);
            }
        }

        private void translateInputs()
        {
            foreach (String str in inputFile.inputs())
            {
                smt2inputs.Add("(declare-const " + str.Trim() + ")");
            }

            foreach (String str in smt2inputs)
            {
                Console.WriteLine(str);
            }
        }

        private void translateInitializations()
        {
            foreach (String str in inputFile.initializations())
            {
                String translStr = "(assert " + str.Trim() + ")";

                foreach (ArrayVariable var in arrayVariables)
                {
                    translStr = var.Replace(translStr, 0);
                }

                smt2initializations.Add(translStr);
            }

            foreach (String str in smt2initializations)
            {
                Console.WriteLine(str);
            }
        }

        private void translateStatement()
        {

            //Console.ReadLine();
                    
            int progLine = 0, smtLine = 0;
            while (progLine < inputFile.statements().Count)
            {
                String str = inputFile.statements().ElementAt(progLine);
                bool branch = false;
                int n, pcIf=0, pcElse=0;

                //Check if line contains a conditional branch (assumption: it's always on the pc)
                if (str.Contains("(ite"))
                {
                    //Calculate length of the loop body
                    branch = true;
                    Smt2Statement stmt = new Smt2Statement(str.Substring(str.IndexOf("(ite")));
                    //Console.WriteLine(String.Join("\r\n", stmt.args));
                    pcIf = convertBv(stmt.args.ElementAt(2));
                    pcElse = convertBv(stmt.args.ElementAt(3));
                    n = pcElse - pcIf;
                }
                else
                {
                    //No loop
                    n = 0;
                }

                //Actions based on whether the current line contains a branch or not
                if (0 == n)
                {
                    //No branch, translate regularly
                    String translStr = "(assert " + str.Trim() + ")";
                    translStr = replaceVars(translStr,smtLine);
                    smt2statements.Add(translStr);
                    smtLine++;
                }
                else
                {
                    //Branch requires to state all possible (pc) states
                    int kmax = searchBound;
                    for (int k = 0; k < kmax; k++)
                    {
                        for (int m = 0; m < n+1; m++)
                        {
                            //Loop condition or body statement
                            str = inputFile.statements().ElementAt(progLine + m);
                            String translStr = "(assert " + str.Trim() + ")";
                            translStr = replaceVars(translStr, k * (n) + m + smtLine);
                            smt2statements.Add(translStr);

                            //Else statement
                            str = inputFile.statements().ElementAt(pcElse);
                            translStr = "(assert " + str.Trim() + ")";
                            translStr = replaceVars(translStr, k * (n) + m + smtLine+1);
                            smt2statements.Add(translStr);
                        }
                        smtLine = smtLine+1;
                    }
                    progLine++;
                    smtLine += kmax*n;
                }
                progLine = progLine + n + 1;
                progSteps = smtLine;
            }
            
            foreach (String str in smt2statements)
            {
                Console.WriteLine(str);
            }
        }

        private String replaceVars(String str, int level)
        {
            String translStr = str;
            foreach (ArrayVariable var in arrayVariables)
            {
                translStr = var.Replace(translStr, level);
            }
            return translStr;
        }

        private void translateConditions()
        {
            foreach (String str in inputFile.conditions())
            {
                String translStr = "(assert (not " + str.Trim() + "))";

                foreach (ArrayVariable var in arrayVariables)
                {
                    translStr = var.Replace(translStr, progSteps);
                }

                smt2conditions.Add(translStr);
            }

            foreach (String str in smt2conditions)
            {
                Console.WriteLine(str);
            }
        }

        public int convertBv(String bv)
        {
            if (bv.Contains("_ bv"))
            {
                int i = bv.IndexOf("_ bv");
                int k = bv.IndexOf(' ', i + 4);
                String str = bv.Substring(i+4,k-(i+4));
                return int.Parse(str);
            }
            return -1;
        }
    }

    class Smt2Statement
    {
        public Smt2Statement(String inStr)
        {
            this.inStr = inStr.Trim();
            this.args = new List<String>();
            disassamble();
        }

        void disassamble()
        {

            String argStr = "";
            int i = 0;
            foreach (char chr in this.inStr)
            {
                if ((chr.Equals(' ') || chr.Equals(')')) && 1 == i)
                {
                    if (argStr.Trim() != String.Empty)
                        args.Add(argStr.Trim());
                    argStr = "";
                    continue;
                }
                else if (chr.Equals('('))
                {
                    i++;
                    if (i < 3)
                    {
                        continue;
                    }
                }
                else if (chr.Equals(')'))
                {
                    i--;
                    if (i < 2)
                    {
                        continue;
                    }
                }
                else if (0 == i)
                {
                    break;
                }
                argStr += chr;
            }
        }

        String inStr
        {
            get;
            set;
        }

        public List<String> args
        {
            get;
            private set;
        }
    }

    class ArrayVariable
    {
        public ArrayVariable(String act, String next, String type, int bitwidth, int cnt)
        {
            this.act = act;
            this.next = next;
            this.type = type;
            this.bitwidth = bitwidth;
            
            //Console.WriteLine("Printing variable: " + this.ToString());
        }

        public String act
        {
            get;
            set;
        }

        public String next
        {
            get;
            set;
        }

        public String type
        {
            get;
            set;
        }

        public int bitwidth
        {
            get;
            set;
        }

        public String Replace(String inStr, int level)
        {
            String outStr = inStr;
            outStr = inStr.Replace(" " + act + " ", " (select " + act + " " + level.ToString() + ") ");
            outStr = outStr.Replace(" " + act + ")", " (select " + act + " " + level.ToString() + "))");
            outStr = outStr.Replace(" " + next + " ", " (select " + act + " " + (level + 1).ToString() + ") ");
            outStr = outStr.Replace(" " + next + ")", " (select " + act + " " + (level + 1).ToString() + "))");
            return outStr;
        }

        public String ToString()
        {
            String returnStr = "([ " + act + " " + next + " ],Type="+type+",BitWidth="+bitwidth.ToString()+")";
            return returnStr;
        }
    }

    class InputFile
    {
        public InputFile(String[] program)
        {
            this.program = new List<String>();
            this.programStr = program;
            if (!initProgram())
                Environment.Exit(0);
        }

        public List<String> program
        {
            get;
            private set;
        }

        private String[] programStr
        {
            get;
            set;
        }

        private int varStart
        {
            get;
            set;
        }

        private int varEnd
        {
            get;
            set;
        }

        private int inputStart
        {
            get;
            set;
        }

        private int inputEnd
        {
            get;
            set;
        }

        private int initStart
        {
            get;
            set;
        }

        private int initEnd
        {
            get;
            set;
        }

        private int stmtStart
        {
            get;
            set;
        }

        private int stmtEnd
        {
            get;
            set;
        }

        private int condStart
        {
            get;
            set;
        }

        private int condEnd
        {
            get;
            set;
        }

        public List<String> variables()
        {
            return program.GetRange(varStart, varEnd - varStart + 1);
        }

        public List<String> inputs()
        {
            return program.GetRange(inputStart, inputEnd - inputStart + 1);
        }

        public List<String> initializations()
        {
            return program.GetRange(initStart, initEnd - initStart + 1);
        }

        public List<String> statements()
        {
            return program.GetRange(stmtStart, stmtEnd - stmtStart + 1);
        }

        public List<String> conditions()
        {
            return program.GetRange(condStart, condEnd - condStart + 1);
        }

        private bool initProgram()
        {
            //Remove comment lines
            foreach (String line in programStr)
            {
                String newLine = line.Trim();

                //Check if line contains more than a newline symbol
                if (newLine.Length > 0)
                {
                    // Check if line does not start with #. If it does, skip it
                    if (!newLine[0].Equals('#'))
                    {
                        // Remove inline comment
                        if (newLine.Contains('#'))
                        {
                            newLine = newLine.Substring(0, newLine.IndexOf('#'));
                        }
                        program.Add(newLine);
                    }
                }
            }

            //Check if program actually has statements
            if (program.Count == 0)
            {
                Console.WriteLine("Program does not contain any statements. Quitting.\n");
                return false;
            }

            //Get args of the first line that declare statement numbers
            String[] progArgs = program.ElementAt(0).Split(' ');

            //Check number of args in the first line (must be 5)
            if (progArgs.Length != 5)
            {
                Console.WriteLine("Number of statement declarations is not 5. Quitting.");
                return false;
            }

            //See if arguments match the declared number
            int nProgArgs = 0;
            int k = 0;
            foreach (String arg in progArgs)
            {
                int progArg = int.Parse(arg);

                switch (k)
                {
                    case 0:
                        varStart = nProgArgs + 1;
                        varEnd = nProgArgs + progArg;
                        break;
                    case 1:
                        inputStart = nProgArgs + 1;
                        inputEnd = nProgArgs + progArg;
                        break;
                    case 2:
                        initStart = nProgArgs + 1;
                        initEnd = nProgArgs + progArg;
                        break;
                    case 3:
                        stmtStart = nProgArgs + 1;
                        stmtEnd = nProgArgs + progArg;
                        break;
                    case 4:
                        condStart = nProgArgs + 1;
                        condEnd = nProgArgs + progArg;
                        break;
                }
                k++;

                nProgArgs += progArg;
            }
            if (program.Count - 1 != nProgArgs)
            {
                Console.WriteLine("Number of program statements ({0}) does not match declared value ({1}). Quitting.\n", program.Count - 1, nProgArgs);
                return false;
            }

            

            return true;
        }

        //Print program statements
        public void printProgram()
        {
            int i = 0;
            Console.WriteLine("Printing program...");
            foreach (String line in program)
            {
                Console.WriteLine(String.Format("{0:000}",i) + ": " + line);
                i++;
            }
        }

        public void printVariables()
        {
            int i = 0;
            Console.WriteLine("Printing variable declarations...");
            foreach (String line in variables())
            {
                Console.WriteLine(String.Format("{0:000}", i + varStart) + ": " + line);
                i++;
            }
        }

        public void printInputs()
        {
            int i = 0;
            Console.WriteLine("Printing inputs declarations...");
            foreach (String line in inputs())
            {
                Console.WriteLine(String.Format("{0:000}", i + inputStart) + ": " + line);
                i++;
            }
        }

        public void printInitializations()
        {
            int i = 0;
            Console.WriteLine("Printing initializations declarations...");
            foreach (String line in initializations())
            {
                Console.WriteLine(String.Format("{0:000}", i + initStart) + ": " + line);
                i++;
            }
        }

        public void printStatements()
        {
            int i = 0;
            Console.WriteLine("Printing statements declarations...");
            foreach (String line in statements())
            {
                Console.WriteLine(String.Format("{0:000}", i + stmtStart) + ": " + line);
                i++;
            }
        }

        public void printConditions()
        {
            int i = 0;
            Console.WriteLine("Printing conditions declarations...");
            foreach (String line in conditions())
            {
                Console.WriteLine(String.Format("{0:000}", i + condStart) + ": " + line);
                i++;
            }
        }
    }
}
